import os.path
from cStringIO import StringIO
from kuai.qmd import *
from kuai.kuaiext import XYZ
from kuai.mol import Atom

class DffQmdFileReader(QuantumDataReader):
    def __init__(self):
        QuantumDataReader.__init__(self)

    def read(self, file):
        result = []
        qmd = QuantumData()
        for i in file:
            if i.startswith("$NUMBER OF ATOMS"):
                qmd.atoms = DffQmdFileReader.read_atoms(file)
            elif i.startswith("$ENERGY"):
                qmd.energy = float(file.next().split()[0])
            elif i.startswith("$FORCES"):
                pass
            elif i.startswith("$FORCE MATRIX"):
                qmd.hessian = DffQmdFileReader.read_hessian(file)
                assert len(qmd.hessian) == len(qmd.atoms) * 3 * (len(qmd.atoms) * 3 + 1) / 2
            elif i.startswith("$$$$"):
                result.append(qmd)
                qmd = QuantumData()
        return result

    @staticmethod
    def read_atoms(file):
        result = None
        for i in file:
            if i.startswith('!'):
                continue
            elif i.startswith('$ATOM'):
                result = []
            elif i.startswith('$END !ATOM'):
                break
            elif result is not None:
                tokens = i.split()
                coords = XYZ(float(tokens[3]), float(tokens[4]), float(tokens[5]))
                result.append(Atom(tokens[1].strip(), coords))
        return result

    @staticmethod
    def read_hessian(file):
        result = []
        for i in file:
            if i.startswith("$END"):
                break
            else:
                result += [float(j.strip()) for j in i.split()]
        return result




reader = DffQmdFileReader()
